#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/sched.h>
#include <linux/tty.h>
#include <linux/version.h>
#include <linux/proc_fs.h>
#include <linux/kernel_stat.h>
#include <asm/uaccess.h>
#include <asm/cputime.h>

#define MODULE_NAME "per_cpu_stats"

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Steven Hofmeyr (LBNL)");
MODULE_DESCRIPTION("Creates per cpu proc files to report CPU stats");

#define __STAT_IDLE 0
#define __STAT_SYS 1
static int stat_type = __STAT_IDLE;
module_param(stat_type, int, S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH);
MODULE_PARM_DESC(stat_type, "Can be 0 for idle (default) or 1 for sys");

// pretty colors for output           
#define _RED_ "\x1b[31m"
#define _GREEN_ "\x1b[32m"
#define _BROWN_ "\x1B[33m"
#define _BLUE_ "\x1b[34m"
#define _MAGENTA_ "\x1b[35m"
#define _CYAN_ "\x1b[36m"
#define _LIGHTGRAY_ "\x1B[37m"
#define _DARKGRAY_ "\x1B[1m\x1B[0m"
#define _LIGHTRED_ "\x1B[1m\x1B[31m"
#define _LIGHTGREEN_ "\x1B[1m\x1B[32m"
#define _YELLOW_ "\x1B[1m\x1B[33m"
#define _LIGHTBLUE_ "\x1B[1m\x1B[34m"
#define _LIGHTMAGENTA_ "\x1B[1m\x1B[35m"
#define _LIGHTCYAN_ "\x1B[1m\x1B[36m"
#define _END_ "\x1b[0m"

#define PRINTCON(fmt, ...) \
  do {                                            \
    char buf[1000];                               \
    snprintf(buf, 1000, fmt, __VA_ARGS__);        \
    printcon(buf);                                \
  } while(0);

static void printcon(char *str) {
  struct tty_struct *my_tty = current->signal->tty;
  if (my_tty != NULL) {
    ((my_tty->driver->ops)->write) (my_tty, str, strlen(str));
    ((my_tty->driver->ops)->write) (my_tty, "\015\012", 2);
  }
}

#ifndef arch_idle_time
#define arch_idle_time(cpu) 0
#endif

static struct proc_dir_entry *pdir;
static int cids[255];

int procfile_read(char *buffer, char **buffer_location,
                  off_t offset, int buffer_length, int *eof, void *data) {
  int ret;
  cputime64_t idle = 0;
  unsigned long long idle_ticks = 0;
  int* cid;

  if (offset > 0) {
    // we have finished read, return 0 
    ret  = 0;
  } else {
    cid = (int*)data;
    if (stat_type == __STAT_IDLE) idle = kstat_cpu(*cid).cpustat.idle;
    else idle = kstat_cpu(*cid).cpustat.system;
    idle_ticks = (unsigned long long)cputime64_to_clock_t(idle);
    //    PRINTCON(_GREEN_ "CPU %d idle time %llu" _END_, *cid, idle_ticks);
    // fill the buffer, return the buffer size 
    ret = sprintf(buffer, "%llu\n", idle_ticks);
    if (ret < 0) {
      printk(KERN_ALERT "Error: could not format idle ticks into output for %s\n", MODULE_NAME);
      return 0;
    } 
  }
  return ret;
}

static int __init init_per_cpu_stats(void) {
  struct proc_dir_entry *pfile;
  char pname[6];
  int i;
  pdir = proc_mkdir(MODULE_NAME, NULL);
  if (pdir == NULL) {
    printk(KERN_ALERT "Error: Could not create /proc/%s\n", MODULE_NAME);
    return -ENOMEM;
  }
  // now create an entry per online cpu
  for_each_online_cpu(i) {
    cids[i] = i;
    sprintf(pname, "cpu%d", i);
    pfile = create_proc_read_entry(pname, 0444, pdir, procfile_read, &cids[i]);
    if (pfile == NULL) {
      remove_proc_entry(MODULE_NAME, NULL);
      printk(KERN_ALERT "Error: Could not create /proc/%s/%s\n", MODULE_NAME, pname);
      return -ENOMEM;
    }
    PRINTCON(_BLUE_ "%s loaded, created /proc/%s/%s" _END_, MODULE_NAME, MODULE_NAME, pname);
  }
  PRINTCON(_BLUE_ "stat type is %d\n" _END_, stat_type);
  return 0;
}

static void __exit cleanup_per_cpu_stats(void) {
  char pname[6];
  int i;
  for_each_online_cpu(i) {
    sprintf(pname, "cpu%d", i);
    remove_proc_entry(pname, pdir);
  }
  remove_proc_entry(MODULE_NAME, NULL);
  PRINTCON(_BLUE_ "%s unloaded" _END_, MODULE_NAME);
}

module_init(init_per_cpu_stats);
module_exit(cleanup_per_cpu_stats);

